from asyncio.subprocess import PIPE
from os import environ
import subprocess
import time

import argparse

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("command", help="an integer to be summed", required=True)
    args = parser.parse_args()
    command = args.command.split("\n")
    i = 0
    while True:
        i += 1
        print(i)
        ret = subprocess.run(
            ["bash", "-c", "./check_gpu_status.sh"], check=True, stdout=PIPE
        )

        for mem_used in map(int, ret.stdout.decode("utf-8").split("\n")[:-1]):
            if mem_used < 500:
                print(
                    "It seems there is a free GPU with memory used: {}".format(mem_used)
                )
                print("try launch my job")
                subprocess.run(command, check=True, env=environ.copy())
                exit(0)
        time.sleep(60)
